#include <sstream>
#include <iomanip>
#include <opencv2/dnn.hpp>

#include "base_util/utils.h"
#include "cls_video.h"

namespace ai {


HBClsVideo::HBClsVideo(HBPackPredictor* model, HBModelManager* manager, LogInfo *lg):
  HBPredictor(model, manager, lg) { }


int HBClsVideo::start_postprocess_thread_imp() {
  int ret = 0;
  std::pair<std::pair<int,int>,int> cur_block;
  // stop_flag2 = common::CameraGrabing;
  // while (stop_flag2 == common::CameraGrabing) {

  //   run_mtx2.lock();
  //   if (post_data2.empty()) {run_mtx2.unlock(); std::this_thread::sleep_for(std::chrono::milliseconds(3));continue;}
  //   cur_block = post_data2.front();
  //   post_data2.pop();
  //   run_mtx2.unlock();

  //   int loop_idx = cur_block.first.second;
  //   std::vector<BaseInfo*> det_infos;
  //   ret = det_hb_->PostProcess(det_infos, cur_block.first.first, this, cur_block.second);  //处理图片
  //   run_mtx3.lock();
  //   out_data3[loop_idx].push(det_infos);  //将处理后的数据返回到det_infos,加入队列中
  //   run_mtx3.unlock();
  // }
  // stop_flag2 = common::CameraOpened;
  return 0;
}

int HBClsVideo::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret;
  // if (blob->imgs.size() != mdl_hb->nchw[0]) {return model_image_batchs_check_error;}
  int batch = 8;

  std::vector<float> im_vec_data;
  int img_buffer_size = mdl_hb->cfg->input_shape[0] * mdl_hb->cfg->input_shape[1] * mdl_hb->cfg->channels;
  // std::cout << "img_buffer_size:" << img_buffer_size << std::endl;
  im_vec_data.resize(batch * img_buffer_size);
  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  for (int idx = 0; idx < batch; idx++) {
    if (!mdl_hb->cfg->transforms->run(blob->imgs[idx], img_blob, mdl_hb->cfg)) { 
      printf("transforms->run fail \n");
      return model_image_channels_check_error;
    }
    // std::cout << "img_blob.img:" << img_blob.img.size() << std::endl;
    // std::memcpy(im_vec_data.data() + img_buffer_size*idx, img_blob.img.data, img_buffer_size* sizeof(float));
    std::memcpy(im_vec_data.data() + img_buffer_size*idx, img_blob.im_vec_data.data(), img_buffer_size* sizeof(float));
  }

  mdl_hb->input_tensors[0].properties.tensorType = HB_DNN_TENSOR_TYPE_F32;
  mdl_hb->input_tensors[0].properties.tensorLayout = HB_DNN_LAYOUT_NCHW;
  mdl_hb->input_tensors[0].properties.validShape.dimensionSize[0] = 8;
  mdl_hb->input_tensors[0].properties.validShape.dimensionSize[1] = 3;
  mdl_hb->input_tensors[0].properties.validShape.dimensionSize[2] = mdl_hb->cfg->input_shape[0];
  mdl_hb->input_tensors[0].properties.validShape.dimensionSize[3] = mdl_hb->cfg->input_shape[1];
  // Align by 16 bytes
  mdl_hb->input_tensors[0].properties.alignedShape.numDimensions = 4;
  mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[0] = 8;
  mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[1] = 3;
  mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[2] = mdl_hb->cfg->input_shape[0];
  mdl_hb->input_tensors[0].properties.alignedShape.dimensionSize[3] = (mdl_hb->cfg->input_shape[1]+16-1) / 16 * 16;
  int ori_length = batch * img_buffer_size * sizeof(float);
  // std::cout << "ori_length:" << ori_length << std::endl;

  hbSysAllocCachedMem(&mdl_hb->input_tensors[0].sysMem[0], ori_length);
  memcpy(mdl_hb->input_tensors[0].sysMem[0].virAddr, im_vec_data.data(), ori_length);
  hbSysFlushMem(&(mdl_hb->input_tensors[0].sysMem[0]), HB_SYS_MEM_CACHE_CLEAN);


  hbDNNTaskHandle_t task_handle = nullptr;
  hbDNNInferCtrlParam infer_ctrl_param;
  HB_DNN_INITIALIZE_INFER_CTRL_PARAM(&infer_ctrl_param);
  hbDNNTensor *output_tensors = mdl_hb->output_tensors.data();
  ret = hbDNNInfer(&task_handle, &output_tensors, mdl_hb->input_tensors.data(), mdl_hb->dnn_handle, &infer_ctrl_param);
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.0.2 hbDNNInfer ret: {}", ret);

  hbDNNWaitTaskDone(task_handle, 0);
  for (int i = 0; i < mdl_hb->output_count; i++) {
    hbSysFlushMem(&mdl_hb->output_tensors[i].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);
  }
  hbDNNReleaseTask(task_handle);    // 释放task handle

  auto data = reinterpret_cast<float *>(mdl_hb->output_tensors[0].sysMem[0].virAddr);
  auto shift = mdl_hb->output_tensors[0].properties.shift.shiftData;
  auto scale = mdl_hb->output_tensors[0].properties.scale.scaleData;
  int n_dim = mdl_hb->output_tensors[0].properties.validShape.numDimensions;
  int *shape = mdl_hb->output_tensors[0].properties.validShape.dimensionSize;
  int output_tensor_size = 1;
  for (int i = 1; i < n_dim; i++) { output_tensor_size *= shape[i]; }

  std::vector<std::pair<float, int>> output_datas(output_tensor_size);
  for (auto i = 0; i < output_tensor_size; i++) {
    float score;
    if (mdl_hb->output_tensors[0].properties.quantiType == SHIFT) {
      // printf("mdl_hb->output_tensors[0].properties.quantiType == SHIFT \n");
      score = static_cast<float>(data[i]) / static_cast<float>(1 << shift[i]);
    } 
    else if (mdl_hb->output_tensors[0].properties.quantiType == SCALE)  {
      // printf("mdl_hb->output_tensors[0].properties.quantiType == SCALE \n");
      score = static_cast<float>(data[i]) * scale[i];
    } 
    else {
      // printf("no quantiType. \n");
      score = static_cast<float>(data[i]);
    }
    // printf("%.4f\n",score);
    output_datas[i] = std::make_pair(score, i); 
  }

  // softmax
  float total_sum = 0;
  for (auto& x : output_datas) {x.first = std::exp(x.first);}
  for (auto& x : output_datas) {total_sum += x.first;}
  for (auto& x : output_datas) {x.first = x.first / total_sum;}

  // top5
  ai::ClsInfo* cls_ifo = new ClsInfo();
  std::sort(output_datas.begin(), output_datas.end(), [](std::pair<float, int>&a, std::pair<float, int>&b){return a.first > b.first;});
  for (int top_idx = 0; top_idx < 5 && top_idx < output_tensor_size; top_idx++) {
    cls_ifo->top5_score.push_back(output_datas[top_idx].first);
    cls_ifo->top5_class.push_back(output_datas[top_idx].second);
  }
  cls_ifo->class_idx = cls_ifo->top5_class[0];
  cls_ifo->score = cls_ifo->top5_score[0];
  det_infos.push_back(cls_ifo);

  return 0;
}




int HBClsVideo::PostProcess(std::vector<BaseInfo*>& det_infos, 
                int cur_block,
                HBModelManager* mng, 
                int md_idx) {

  return 0;
}



}  // namespace hb
